import streamlit as st
import torch
from jinja2 import Template

from databases.milvus import MilvusDB
from databases.mysql import MySQLDB
from models.chat.load_chat_model import LoadChatModel

import plotly.express as px
import pandas as pd

# 设置页面配置
st.set_page_config(
    page_title="Streamlit Simple Demo",  # 设置页面标题
    page_icon=":robot:",  # 设置页面图标
    layout="wide"  # 设置页面布局为宽屏
)

# 在侧边栏添加滑块用于调整模型参数
max_length = st.sidebar.slider("max_length", 0, 32768, 8192, step=1)
top_p = st.sidebar.slider("top_p", 0.00, 1.00, 0.01, step=0.01)
temperature = st.sidebar.slider("temperature", 0.00, 1.00, 0.01, step=0.01)

# 检查是否存在名为"history"的键
if "history" not in st.session_state:
    # 如果不存在，则将"history"初始化为空列表
    st.session_state.history = []

# 添加清理会话历史的按钮
buttonClean = st.sidebar.button("清理会话历史", key="clean")
if buttonClean:
    # 清空会话历史和过去的键值
    st.session_state.history = []
    st.session_state.past_key_values = None
    # 如果有GPU，则清空CUDA缓存
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

# 遍历并显示会话历史中的消息
for i, message in enumerate(st.session_state.history):
    if message["role"] == "user":
        # 如果消息是用户的，则使用用户头像和样式显示
        with st.chat_message(name="user", avatar="user"):
            st.markdown(message["content"])
    else:
        # 如果消息是助手的，则使用助手头像和样式显示
        with st.chat_message(name="assistant", avatar="assistant"):
            st.markdown(message["content"])

# 添加用户输入的placeholder
with st.chat_message(name="user", avatar="user"):
    input_placeholder = st.empty()

# 显示 SQL 查询结果
with st.chat_message(name="tool", avatar="ai"):
    tool_placeholder = st.empty()

# 添加助手响应的placeholder
with st.chat_message(name="assistant", avatar="assistant"):
    assistant_placeholder = st.empty()


# ==========================以上是WebUI================================================================================

@st.cache_resource
def init():
    _mysqldb = MySQLDB()
    _milvusdb = MilvusDB()
    _model = LoadChatModel()
    _model.load()
    return _model, _mysqldb, _milvusdb


# 加载Chatglm3, mysqldb, milvusdb
model, mysqldb, milvusdb = init()


# 拼接提示模板
def get_prompt_text(ddl, sql, sql_describe, question):
    # 创建 Jinja2 模板
    template_string = """用户提供一个问题，而您提供SQL。您将只使用SQL代码而不使用任何解释进行响应。仅使用SQL代码进行响应。不要回答任何解释——只回答代码。

您可以使用以下DDL语句作为可用表的参考。还可以使用对过去问题的回答来指导您：
{{ ddl }}

您可以使用以下文档作为可用表格的参考。还可以使用对过去问题的回答来指导您：
{{ sql_describe }}

您可以使用以下SQL语句作为可用表的参考。还可以使用对过去问题的回答来指导您：
{{ sql }}

用户问题:
{{ question }}"""
    template = Template(template_string)

    # 提供变量值并渲染模板
    return template.render(ddl=ddl, sql=sql, sql_describe=sql_describe, question=question)


# 渲染图表
def rendering_charts(chart_data):
    # 将数据转换为DataFrame
    df = pd.DataFrame(chart_data)

    # 动态获取所有键
    all_keys = df.columns

    # 允许用户选择X轴和Y轴的键
    x_column = st.selectbox('选择X轴键', all_keys, key='x_column')
    y_column = st.selectbox('选择Y轴键', all_keys, key='y_column', index=1)  # 通过index参数设置默认选项

    # 创建图表
    fig = px.bar(df, x=x_column, y=y_column, title='', labels={x_column: x_column, y_column: y_column})

    # 渲染图表
    st.plotly_chart(fig)


# 用户输入问题
question = st.chat_input("请输入您的问题")

# 如果有问题，则进行处理
if question:
    try:
        # 显示用户输入
        input_placeholder.markdown(question)

        # 返回业务SQL
        result = milvusdb.search_milvus(question)[0]

        # 拼接提示模板
        prompt_text = get_prompt_text(ddl=result.get('ddl'), sql=result.get('sql'), sql_describe=result.get('sql_describe'), question=question)
        # 模型推理,填充SQL
        response = model.chat(prompt_text)

        # 将SQL返回到前端页面
        tool_placeholder.markdown(response[0])

        # 使用SQl查询 MySQL
        result = mysqldb.query(response[0])
        # 将结果数据返回到前端页面
        assistant_placeholder.markdown(result)

        # 渲染图表
        rendering_charts(result)
    except Exception as e:
        assistant_placeholder.markdown(e)
